# 05b_predict_lstm.py
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 减少TensorFlow日志输出
import sys
import pandas as pd
import numpy as np
from sklearn.preprocessing import RobustScaler
from tensorflow.keras.models import load_model
import tensorflow as tf

# --- 智能GPU配置 ---
def configure_gpu(force_cpu=False):
    """配置GPU设备，支持CPU回退"""
    if force_cpu:
        # 强制使用CPU
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        print("ℹ️ 强制使用CPU模式")
        return False
        
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # 允许GPU内存增长
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            print(f"✅ GPU配置成功！将使用GPU加速")
            return True
        except Exception as e:
            print(f"⚠️ GPU配置失败: {e}")
            print("   自动回退到CPU模式")
            os.environ['CUDA_VISIBLE_DEVICES'] = ''
            return False
    else:
        print("ℹ️ 未检测到GPU设备，使用CPU")
        return False

# 检查环境变量，允许用户选择CPU模式
force_cpu_mode = os.getenv('FORCE_CPU', 'false').lower() == 'true'
gpu_available = configure_gpu(force_cpu=force_cpu_mode)
import efinance as ef
import talib
import joblib
from pypinyin import pinyin, Style

# --- 核心参数 ---
TIME_STEPS = 60  # 与训练时相同的时间步长

def get_company_name(stock_code):
    """获取股票中文名称并转换为拼音"""
    try:
        info = ef.stock.get_base_info(stock_code)
        name = info['股票名称']
        pinyin_name = "".join(sum(pinyin(name, style=Style.NORMAL), []))
        return pinyin_name
    except Exception:
        return ""

def calculate_indicators(df):
    """为DataFrame计算所需的技术指标 (输入df应有英文列名)"""
    close = df['close'].astype(float)
    high = df['high'].astype(float)
    low = df['low'].astype(float)
    volume = df['volume'].astype(float)
    
    df['sma5'] = talib.SMA(close, timeperiod=5)
    df['sma10'] = talib.SMA(close, timeperiod=10)
    df['sma20'] = talib.SMA(close, timeperiod=20)
    df['ema12'] = talib.EMA(close, timeperiod=12)
    df['ema26'] = talib.EMA(close, timeperiod=26)
    macd, macdsignal, macdhist = talib.MACD(close, fastperiod=12, slowperiod=26, signalperiod=9)
    df['macd'] = macd
    df['macdsignal'] = macdsignal
    df['macdhist'] = macdhist
    df['rsi'] = talib.RSI(close, timeperiod=14)
    upper, middle, lower = talib.BBANDS(close, timeperiod=20)
    df['upperband'] = upper
    df['middleband'] = middle
    df['lowerband'] = lower
    df['atr'] = talib.ATR(high, low, close, timeperiod=14)
    df['adx'] = talib.ADX(high, low, close, timeperiod=14)
    
    return df

def predict_next_day(stock_code, frequency):
    """
    使用已训练的LSTM模型预测下一个周期的收盘价
    """
    model_filename = f'output/models/lstm_model_{stock_code}_{frequency}.h5'
    X_scaler_filename = f'output/scalers/X_scaler_{stock_code}_{frequency}.pkl'
    y_scaler_filename = f'output/scalers/y_scaler_{stock_code}_{frequency}.pkl'
    data_with_indicators_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'

    # 检查必要文件是否存在
    if not os.path.exists(model_filename):
        print(f"错误: 未找到模型文件 '{model_filename}'。请先运行 'train_lstm' 命令进行训练。")
        sys.exit(1)
    
    if not os.path.exists(X_scaler_filename):
        print(f"错误: 未找到特征归一化器文件 '{X_scaler_filename}'。请重新训练模型。")
        sys.exit(1)
        
    if not os.path.exists(y_scaler_filename):
        print(f"错误: 未找到目标归一化器文件 '{y_scaler_filename}'。请重新训练模型。")
        sys.exit(1)

    print(f"正在加载LSTM模型: {model_filename}...")
    model = load_model(model_filename)
    
    print(f"加载训练时的归一化器...")
    X_scaler = joblib.load(X_scaler_filename)
    y_scaler = joblib.load(y_scaler_filename)
    
    # 获取特征列表（从历史数据文件中读取列名）
    if not os.path.exists(data_with_indicators_filename):
        print(f"错误: 找不到历史数据文件 '{data_with_indicators_filename}'。")
        sys.exit(1)
        
    hist_df = pd.read_csv(data_with_indicators_filename)
    features = [col for col in hist_df.columns if col not in ['date', 'code', 'target', 'open', 'close']]
    
    print("正在获取最新的市场数据 (最近200条以确保指标稳定)...")
    latest_data_raw = ef.stock.get_quote_history(stock_code, klt=101, limit=200)
    
    # --- 数据处理流程与训练时保持完全一致 ---
    columns_to_keep = ['日期', '股票代码', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
    latest_df = latest_data_raw[columns_to_keep].copy()
    
    rename_map = {
        '日期': 'date', '股票代码': 'code', '开盘': 'open', '收盘': 'close', 
        '最高': 'high', '最低': 'low', '成交量': 'volume', '成交额': 'amount', 
        '振幅': 'amplitude', '涨跌幅': 'change_percent', '涨跌额': 'change_amount', '换手率': 'turnover'
    }
    latest_df.rename(columns=rename_map, inplace=True)

    try:
        latest_df = calculate_indicators(latest_df)
    except KeyError as e:
        print(f"在 calculate_indicators 中发生 KeyError: {e}")
        print("latest_df 的列名:", latest_df.columns)
        sys.exit(1)
    
    if len(latest_df) < TIME_STEPS:
        print(f"错误: 获取到的最新数据不足 {TIME_STEPS} 条，无法进行预测。")
        sys.exit(1)
        
    last_sequence_df = latest_df.tail(TIME_STEPS)
    last_sequence_df.replace([np.inf, -np.inf], np.nan, inplace=True)
    if last_sequence_df.isnull().values.any():
        print("警告: 最新数据序列中存在NaN值，将使用前向填充处理。")
        last_sequence_df.fillna(method='ffill', inplace=True)
        last_sequence_df.fillna(method='bfill', inplace=True)

    last_sequence_df_aligned = last_sequence_df[features]
    scaled_sequence = X_scaler.transform(last_sequence_df_aligned)
    input_data = np.reshape(scaled_sequence, (1, TIME_STEPS, len(features)))

    # 模型预测（输出是归一化后的值）
    prediction_scaled = model.predict(input_data)[0]
    # 关键修复：使用目标归一化器将预测结果反归一化回真实价格
    prediction_original = y_scaler.inverse_transform(prediction_scaled.reshape(1, -1))[0]
    predicted_open = prediction_original[0]
    predicted_close = prediction_original[1]

    company_name = get_company_name(stock_code)
    last_close = latest_df['close'].iloc[-1]
    change = predicted_close - last_close
    change_percent = (change / last_close) * 100
    
    print("\n--- LSTM 双目标模型预测结果 ---")
    print(f"公司名称: {company_name} ({stock_code})")
    print(f"当前周期: {frequency}")
    print(f"最新收盘价: {last_close:.2f}")
    print(f"预测下一周期开盘价: {predicted_open:.2f}")
    print(f"预测下一周期收盘价: {predicted_close:.2f}")
    print(f"预测涨跌 (基于收盘价): {change:.2f} ({change_percent:+.2f}%)")
    if change > 0:
        print("预测结果: 看涨 📈")
    else:
        print("预测结果: 看跌 📉")
    print("------------------------------\n")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("用法: python3 05b_predict_lstm.py <股票代码> <frequency>")
        sys.exit(1)
    
    stock_code = sys.argv[1]
    frequency = sys.argv[2]
    predict_next_day(stock_code, frequency)